import os
from torch.utils.data import Dataset
from utils import *
from torchvision import transforms
transform=transforms.Compose([
    transforms.ToTensor()
])

class Mydataset(Dataset):
    def __init__(self,path):
        self.path=path
        self.name=os.listdir(os.path.join(path,'SegmentationClass'))
    def __len__(self):
        return len(self.name)
    def __getitem__(self, index):
        segmen_name=self.name[index]   ### xx.png
        segmen_path=os.path.join(self.path,'SegmentationClass',segmen_name)
        image_path=os.path.join(self.path,'JPEGImages',segmen_name.replace('png','jpg'))
        segmen_image=keep_img_size(segmen_path)
        image=keep_img_size(image_path)
        return transform(segmen_image),transform(image)
if __name__ == '__main__':
    path='D:\datatrain\VOCdevkit\VOC2007'
    data=Mydataset(path)
    print(data[0][0].shape)
    print(data[0][1].shape)

